import numpy as np
from scipy.fft import fft
from scipy import signal
import pywt

def read_file(file_path):
    """
    读取文本文件中的振动加速度信号，并转换成数字列表。
    """
    with open(file_path, 'r') as file:
        content = file.readlines()
    data = [float(x.strip()) for x in content]
    return data

def get_fft(data, sample_rate):
    """
    将输入信号转换到频域，并返回频域振幅值的规范化结果。
    """
    n = len(data)
    freqs = np.fft.fftfreq(n, 1/sample_rate)
    mask = (freqs > 0) & (freqs < sample_rate/2)
    fft_data = fft(data)
    fft_abs = np.abs(fft_data[mask])
    fft_norm = fft_abs / np.max(fft_abs)
    return fft_norm

def remove_low_freq(data, fs=10000, lowcut=10):
    """
    将输入信号进行低通滤波，去除低于设定频率的部分。
    """
    nyquist = 0.5 * fs
    low = lowcut / nyquist
    b, a = signal.butter(5, low, btype='low')
    filtered_data = signal.filtfilt(b, a, data)
    return filtered_data

def get_similarity(file_path1, file_path2, num_freqs=17000, time_weight=0.6):
    """
    计算两个文本文件中振动加速度信号的相似度。
    """
    # 读取数据并进行低通滤波
    data1 = read_file(file_path1)
    data2 = read_file(file_path2)
    filtered_data1 = remove_low_freq(data1)
    filtered_data2 = remove_low_freq(data2)
    # 小波分解
    cA1, cD1 = pywt.dwt(filtered_data1, 'db4')
    cA2, cD2 = pywt.dwt(filtered_data2, 'db4')
    # 计算时域距离
    distance1 = np.linalg.norm(np.array(cA1) - np.array(cA2))
    # 计算频域距离
    fft_norm1 = get_fft(cD1,10000)
    fft_norm2 = get_fft(cD2,10000)
    # 提取前N个主要频率的振幅值
    freq_idx1 = np.argsort(fft_norm1)[::-1][:num_freqs]
    freq_idx2 = np.argsort(fft_norm2)[::-1][:num_freqs]
    freq_amp1 = fft_norm1[freq_idx1]
    freq_amp2 = fft_norm2[freq_idx2]
    # 归一化处理
    freq_amp1 = freq_amp1 / np.max(freq_amp1)
    freq_amp2 = freq_amp2 / np.max(freq_amp2)
    # 计算余弦相似度
    freq_cos_similarity = np.dot(freq_amp1, freq_amp2) / (np.linalg.norm(freq_amp1) * np.linalg.norm(freq_amp2))
    # 计算频域误差
    freq_error = np.linalg.norm(fft_norm1 - fft_norm2)
    # 计算综合距离
    distance = np.sqrt((time_weight * distance1) ** 2 + ((1-time_weight) * (1-freq_cos_similarity )) ** 2 + freq_error)
    # 计算相似度
    similarity = 1 / (1 + distance)
    return similarity


# 计算每对文件之间的相似度
file_paths = ['7.txt', '8.txt', '2.txt']
num_files = len(file_paths)
similarities = np.zeros((num_files, num_files))
for i in range(num_files):
    for j in range(i + 1, num_files):
        similarity = get_similarity(file_paths[i], file_paths[j])
        similarities[i, j] = similarity
        similarities[j, i] = similarity


# 将相似度矩阵中的每个元素平方后进行加权，并求和
weighted_sum = np.sum(np.square(similarities))  # 每个元素平方后求和
mean_similarity = np.sqrt(weighted_sum / (num_files * (num_files - 1) / 2))  # 平均值的平方根

print(f"Mean similarity: {mean_similarity}")
print("Similarities:")
print(similarities)